import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from pydantic.v1 import BaseModel, Field

load_dotenv()


# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)


# class Classification(BaseModel):
#     """
#     定义一个数据模型，根据该类型完成文本的分类
#     """
#     # 文本的情感倾向，预期为字符串类型
#     sentiment: str = Field(description="文本的情感")
#
#     # 文本的攻击性，预期为1-10的整数
#     aggressiveness: int = Field(description="描述文本的攻击性，数字越大表示越有攻击性")
#
#     # 文本使用的语言
#     language: str = Field(description="文本使用的语言")

# pydantic v1版本的才能正常解析输出？why？
class Classification(BaseModel):
    """
    定义一个数据模型，根据该类型完成文本的分类
    """
    # 文本的情感倾向，预期为字符串类型
    sentiment: str = Field(..., enum=['happy', 'neutral', 'sad'], description="文本的情感")

    # 文本的攻击性，预期为1-10的整数
    aggressiveness: int = Field(..., enum=[1, 2, 3, 4, 5], description="描述文本的攻击性，数字越大表示越有攻击性")

    # 文本使用的语言
    language: str = Field(description="文本使用的语言")


tagging_prompt = ChatPromptTemplate([
    """
    从以下段落的未结构化文本中提取相关信息。
    只提取'Classification'类中提到的相关的属性。
    段落：
    {input}
    """
])

# tagging_prompt = ChatPromptTemplate(
#     messages=[
#         HumanMessagePromptTemplate.from_template(
#             """
#             从以下段落的未结构化文本中提取相关信息。
#             只提取'Classification'类中提到的相关的属性。
#             段落：
#             {input}
#             """
#         )
#     ],
#     input_variables=["input"]
# )

chain = tagging_prompt | model.with_structured_output(Classification)

input_text = "They are comfy and can with stand everything and don't break or bend and cheap they fit perfectly and " \
             "light i can run in them there durable and love the look. "
input_text = "这个商品质量太差了，质量不好，颜色和图片的不一致。"
# result: Classification = chain.invoke({'input': input_text})
result: Classification = chain.invoke(input_text)
print(result)
